"""
数据传输服务模块
负责样本数据的传输和归集功能
"""

import uuid
import json
import shutil
import zipfile
from pathlib import Path
from typing import Dict
from datetime import datetime
from sqlalchemy.orm import Session

from database import Dataset
from config import config
from exceptions import DatasetException, ValidationException, ErrorCode
from utils import LogManager, DatasetUtils

# 初始化日志系统
LogManager.initialize_logging(config.log.level)
logger = LogManager.get_logger(__name__)


class DataTransferService:
    """数据传输服务"""
    
    def __init__(self):
        # 使用统一的配置管理
        self.config = config
        self.logger = LogManager.get_logger(__name__)
        
        # 确保目录存在
        self.config.ensure_directories()
    
    async def transfer_samples(self, source_path: str, file_format: str = "zip") -> Dict:
        """传输样本数据"""
        try:
            source = Path(source_path)
            if not source.exists():
                raise DatasetException(
                    f"源路径不存在: {source_path}",
                    ErrorCode.DATASET_NOT_FOUND,
                    details={"source_path": source_path}
                )
            
            # 生成目标路径
            transfer_id = str(uuid.uuid4())
            target_dir = self.config.get_temp_path(f"transfer_{transfer_id}")
            target_dir.mkdir(parents=True, exist_ok=True)
            
            self.logger.info(f"开始传输样本: {source_path} -> {target_dir}")
            
            if file_format.lower() == "zip":
                await self._transfer_zip_file(source, target_dir)
            else:
                await self._transfer_directory(source, target_dir)
            
            # 验证传输结果
            if not target_dir.exists() or not any(target_dir.iterdir()):
                raise DatasetException(
                    "数据传输失败，目标目录为空",
                    ErrorCode.DATASET_EMPTY,
                    details={"target_dir": str(target_dir)}
                )
            
            self.logger.info(f"样本传输完成: {transfer_id}")
            
            return {
                "transfer_id": transfer_id,
                "status": "completed",
                "target_path": str(target_dir),
                "file_count": len(list(target_dir.rglob("*")))
            }
            
        except Exception as e:
            self.logger.error(f"样本传输失败: {e}")
            raise
    
    async def collect_samples(self, db: Session, collection_data: dict) -> Dict:
        """归集样本数据到数据集中"""
        try:
            # 验证输入参数
            required_fields = ["dataset_name", "source_paths"]
            for field in required_fields:
                if field not in collection_data:
                    raise ValidationException(
                        f"缺少必需字段: {field}",
                        details={"required_fields": required_fields}
                    )
            
            dataset_name = collection_data["dataset_name"]
            source_paths = collection_data["source_paths"]
            
            # 创建目标数据集目录
            dataset_path = self.config.get_dataset_path(dataset_name)
            images_dir = dataset_path / "images"
            labels_dir = dataset_path / "labels"
            
            for directory in [images_dir, labels_dir]:
                directory.mkdir(parents=True, exist_ok=True)
            
            # 归集文件
            collected_files = 0
            for source_path in source_paths:
                collected_files += await self._collect_from_source(
                    Path(source_path), images_dir, labels_dir
                )
            
            if collected_files == 0:
                raise DatasetException(
                    "没有找到有效的样本文件",
                    ErrorCode.DATASET_EMPTY,
                    details={"source_paths": source_paths}
                )
            
            # 分析数据集
            label_files = list(labels_dir.glob("*.txt"))
            class_distribution = DatasetUtils.analyze_class_distribution(label_files)
            
            # 创建数据集信息记录
            self._create_dataset_info(db, dataset_name, dataset_path, class_distribution, collected_files)
            
            self.logger.info(f"样本归集完成: {dataset_name}, 文件数: {collected_files}")
            
            return {
                "dataset_name": dataset_name,
                "status": "completed",
                "dataset_path": str(dataset_path),
                "file_count": collected_files,
                "class_distribution": class_distribution
            }
            
        except Exception as e:
            self.logger.error(f"样本归集失败: {e}")
            raise
    
    async def _transfer_zip_file(self, source: Path, target_dir: Path):
        """传输ZIP文件"""
        try:
            with zipfile.ZipFile(source, 'r') as zip_ref:
                zip_ref.extractall(target_dir)
        except zipfile.BadZipFile:
            raise DatasetException(
                f"无效的ZIP文件: {source}",
                ErrorCode.DATASET_INVALID_FORMAT,
                details={"source": str(source)}
            )
    
    async def _transfer_directory(self, source: Path, target_dir: Path):
        """传输目录"""
        try:
            shutil.copytree(source, target_dir, dirs_exist_ok=True)
        except Exception as e:
            raise DatasetException(
                f"目录传输失败: {e}",
                ErrorCode.DATASET_CORRUPTED,
                details={"source": str(source), "target": str(target_dir)}
            )
    
    async def _collect_from_source(self, source: Path, images_dir: Path, labels_dir: Path) -> int:
        """从源路径归集文件"""
        collected = 0
        
        if source.is_file():
            # 单个文件
            if source.suffix.lower() in ['.jpg', '.jpeg', '.png', '.bmp']:
                shutil.copy2(source, images_dir / source.name)
                collected += 1
            elif source.suffix.lower() == '.txt':
                shutil.copy2(source, labels_dir / source.name)
                collected += 1
        else:
            # 目录
            for file_path in source.rglob("*"):
                if file_path.is_file():
                    if file_path.suffix.lower() in ['.jpg', '.jpeg', '.png', '.bmp']:
                        shutil.copy2(file_path, images_dir / file_path.name)
                        collected += 1
                    elif file_path.suffix.lower() == '.txt' and file_path.parent.name == 'labels':
                        shutil.copy2(file_path, labels_dir / file_path.name)
                        collected += 1
        
        return collected
    
    def _create_dataset_info(self, db: Session, dataset_name: str, dataset_path: Path, 
                           class_distribution: Dict, file_count: int):
        """创建数据集信息记录"""
        try:
            dataset = Dataset(
                dataset_id=str(uuid.uuid4()),
                dataset_name=dataset_name,
                name=dataset_name,
                dataset_path=str(dataset_path),
                dataset_type="collected",
                class_count=len(class_distribution),
                image_count=file_count,
                sample_count=file_count,
                class_distribution=json.dumps(class_distribution),
                data_format="IMAGE",
                uploader="system",
                data_size="Unknown",
                upload_time=datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                time_range="Unknown",
                update_frequency="Once"
            )
            
            db.add(dataset)
            db.commit()
            
        except Exception as e:
            self.logger.warning(f"创建数据集信息记录失败: {e}") 